docs(backend): document backend support and workflows - #5929

Actionable comments posted: 1

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@doc/train/training.md`:
- Line 23: Remove the shell prompt marker from each command-only example, or add
corresponding output, to eliminate MD014 warnings: update the TensorFlow 2,
PyTorch, and PyTorch-Exportable examples at doc/train/training.md lines 23, 30,
and 37, and the TensorFlow 2 and PyTorch-Exportable examples at
doc/freeze/freeze.md lines 21 and 52.

Pull request overview

This PR adds user-facing documentation for selecting and using the TensorFlow 2 (eager) and PyTorch-Exportable (pt-expt) backends, covering backend selection plus train, freeze, and compress CLI usage across the relevant docs pages.

Changes:

• Add TF2 and PyTorch-Exportable tabs to the training instructions.
• Add TF2 and PyTorch-Exportable tabs to the freeze and compress docs (including --lower-kind mention for pt-expt).
• Extend doc/backend.md to include the two backends, their model/checkpoint formats, and inference suffix detection notes.

Reviewed changes

Copilot reviewed 4 out of 4 changed files in this pull request and generated 3 comments.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

Actionable comments posted: 1

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@doc/model/train-fitting-dos.md`:
- Around line 4-11: Update the DOS training documentation to cover TensorFlow 2
and PyTorch-Exportable routes alongside the existing dp --tf train and dp --pt
train examples, adding dp --tf2 train and dp --pt-expt train tabs. If either
backend cannot train DOS models, instead revise the supported-backends
description to explicitly mark it as conversion/inference-only, consistent with
JAX.

---

Outside diff comments:
In `@doc/backend.md`:
- Around line 114-117: Update the `.pt2` inference guidance in the backend
documentation to distinguish DPA4 `dp --pt freeze` artifacts from
PyTorch-Exportable/AOTInductor artifacts, explicitly identifying which ABI each
`.pt2` file uses. Align the wording with the separate PT and pt_expt code paths
so users do not load a `.pt2` graph through the wrong backend contract.

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

This is a large and mostly accurate sweep -- I checked the great majority of the new support claims against backend registration at head and they hold up, including the ones that are easy to get wrong. Specifically confirmed correct: every dp --tf2 / dp --pt-expt / dp --jax flag against the aliases in deepmd/backend/*.py; the dp --pt_expt to dp --pt-expt and .xlo to .hlo corrections (both were real doc bugs); all the frozen-model suffixes against each backend's suffixes; dp --jax freeze --hessian; the JAX energy-loss-only and no---init-frz-model caveats; the TF2 and Paddle additions on dpa2/dpa3/se_e3_tebd; and the DPA4 note, where JAX really does have deepmd/jax/descriptor/dpa4.py but no jax/model/dpa4_model.py, so "descriptor-only" is exactly right.

Two claims do not hold. Both are inline.

One smaller thing I did not raise inline. The Hessian rewording in overall.md and train-energy-hessian.md is not wrong -- it names dp --pt freeze and the TorchScript route specifically -- but hessian_mode is honored only at deepmd/pt/model/model/__init__.py#L371-L374, and neither pt_expt/model/get_model.py nor dpmodel/model/model.py mention it, so PyTorch-Exportable cannot construct a Hessian model at all. Since this PR establishes that {{ pytorch_icon }} covers both, the page is now silent on pt_expt where it used to be blanket. Worth a clause if you touch these lines again; not worth a round trip on its own.

The current DPA4 native-spin multi-rank guidance contradicts both the implementation and itself. The graph-route/native-spin section of doc/model/dpa4.md says that native-spin .pt2 archives carry no with-comm artifact, have has_comm_artifact=false, and are single-rank only. At this head, _needs_with_comm_artifact explicitly admits a graph-lowered NativeSpinModelKind whose descriptor communicates across ranks; test_native_spin_needs_with_comm_on_the_graph_lower_only asserts that behavior; and DeepSpinPTExpt loads and dispatches forward_lower_with_comm.pt2 for multi-rank graph native spin. The same page also correctly states near the earlier multi-rank overview and the Limitations section that pt_expt native spin supports multi-rank. Please remove or update the obsolete single-rank paragraphs so this page documents one contract that matches the current export and C++ runtime.

The documented PyTorch-Exportable compression workflow does not actually put a compressed model into the executable artifact for the ordinary .pte example. Except for the special fused DPA1 graph .pt2 branch, enable_compression passes the original uncompressed model dictionary to deserialize_to_file and stores the compressed dictionary only as model_json_override. DeepEval reconstructs that model.json object for metadata and introspection, but energy inference runs exported_module or _pt2_runner, so the tabulated model is not executed. I reproduced the documented se_e2_a-style .pte workflow locally: the input and output exported graphs were identical (711 nodes each), the output executable contained no tabulate_fusion operator, while only model.json contained the compression state. Therefore model-compress.pte provides no compression speed or memory benefit even though this page and backend.md present pt_expt compression as generally supported. Please either export the genuinely compressed executable or document the narrow supported case instead of advertising the ordinary .pte workflow as compressed.

The newly documented JAX .jax frozen-model route is not usable by the normal DeePMD inference/test interface. Please correct the format contract or implement the missing loader support.

The newly documented TensorFlow 2 compression workflow fails for the normalized default se_e2_a configuration. The implementation or the documented support contract needs to be corrected.

Both points are addressed, and I checked the files at 3bc4ff06a rather than taking the replies at their word.

On the pairtab page: Paddle is gone from the title and from both notes, and DP is back on the linear-combination note and on linear.md. Since the fix introduced a claim I had not checked in the first round, I verified it too — the new "PyTorch-Exportable" on both lines holds, because deepmd/pt_expt/model/ carries both dp_zbl_model.py and dp_linear_model.py. The asymmetry between the two notes is also right rather than an oversight: JAX and TF2 appear on the short-range interpolation line, where each has pairtab_atomic_model.py and dp_zbl_model.py, but not on the linear-combination line, where neither has a dp_linear_model.py — only linear_atomic_model.py. And Paddle genuinely has none of them.

On the compression page: switching the headline example to -i frozen_model.jax -o compressed_model.jax removes the case that would have raised. The prose now points at .jax as the general lossless path and defers the StableHLO question to the descriptor pages; I confirmed that deferral is not a dead pointer, since train-se-e2-a.md and train-se-atten.md both state the type_one_side requirement for .hlo export. One small thing for the record only: your reply said the type_one_side requirement was documented for compressed .hlo export, and in compress.md itself it is referred out rather than stated. That is the better placement given the requirement is descriptor-specific, so I am not asking for a change — just noting the wording difference in case you expected it inline.

The Hessian point from my earlier review body is resolved as well. train-energy-hessian.md now says "Supported backends: PyTorch-TorchScript, JAX" and labels the tabs the same way, instead of the blanket PyTorch icon. That matches where hessian_mode is actually honored — deepmd/pt/, deepmd/jax/, and deepmd/dpmodel/model/make_hessian_model.py, with nothing on the pt_expt side — so the page no longer implies coverage that does not exist.

Both points are addressed, and I checked the files at 3bc4ff06a rather than taking the replies at their word.

On the pairtab page: Paddle is gone from the title and from both notes, and DP is back on the linear-combination note and on linear.md. Since the fix introduced a claim I had not checked in the first round, I verified it too — the new "PyTorch-Exportable" on both lines holds, because deepmd/pt_expt/model/ carries both dp_zbl_model.py and dp_linear_model.py. The asymmetry between the two notes is also right rather than an oversight: JAX and TF2 appear on the short-range interpolation line, where each has pairtab_atomic_model.py and dp_zbl_model.py, but not on the linear-combination line, where neither has a dp_linear_model.py — only linear_atomic_model.py. And Paddle genuinely has none of them.

On the compression page: switching the headline example to -i frozen_model.jax -o compressed_model.jax removes the case that would have raised. The prose now points at .jax as the general lossless path and defers the StableHLO question to the descriptor pages; I confirmed that deferral is not a dead pointer, since train-se-e2-a.md and train-se-atten.md both state the type_one_side requirement for .hlo export. One small thing for the record only: your reply said the type_one_side requirement was documented for compressed .hlo export, and in compress.md itself it is referred out rather than stated. That is the better placement given the requirement is descriptor-specific, so I am not asking for a change — just noting the wording difference in case you expected it inline.

The Hessian point from my earlier review body is resolved as well. train-energy-hessian.md now says "Supported backends: PyTorch-TorchScript, JAX" and labels the tabs the same way, instead of the blanket PyTorch icon. That matches where hessian_mode is actually honored — deepmd/pt/, deepmd/jax/, and deepmd/dpmodel/model/make_hessian_model.py, with nothing on the pt_expt side — so the page no longer implies coverage that does not exist.

Pull request overview

Copilot reviewed 33 out of 33 changed files in this pull request and generated no new comments.

## PyTorch-TorchScript Implementation {{ pytorch_icon }}

Currently, parallel training in pytorch version is implemented in the form of PyTorch Distributed Data Parallelism [DDP](https://pytorch.org/docs/stable/generated/torch.nn.parallel.DistributedDataParallel.html).
DeePMD-kit will decide whether to launch the training in parallel (distributed) mode or in serial mode depending on your execution command.

In the TensorFlow backend, {ref}`env_protection <model[standard]/descriptor[se_e2_r]/env_protection>` cannot be set to a non-zero value.
In the PyTorch-TorchScript, JAX, and DP backends, {ref}`type_one_side <model[standard]/descriptor[se_e2_r]/type_one_side>` cannot be set to `false`.

Re-reviewed the current head after the DPA4/SeZM follow-up. The documentation now correctly separates the dp --pt training route from the AOTInductor .pt2 export/runtime path and explicitly states that DPA4/SeZM does not use ordinary TorchScript. The prior backend-contract concerns are resolved, and I found no new blocking issue.

One remaining backend-classification error is noted inline.

One additional backend-workflow omission is noted inline.

One additional PyTorch-Exportable training-workflow omission is noted inline.

One additional graph-freeze behavior mismatch is noted inline.

One additional backend-workflow documentation gap is noted inline.

One additional spin-scheme semantic mismatch is noted inline.

Pull request overview

Copilot reviewed 34 out of 34 changed files in this pull request and generated no new comments.

DP compression accepts native `.dp` and `.yaml` models.

Pull request overview

Copilot reviewed 34 out of 34 changed files in this pull request and generated no new comments.

> **Supported backends**: TensorFlow {{ tensorflow_icon }}, TensorFlow 2
> {{ tensorflow_icon }}, PyTorch {{ pytorch_icon }}, Paddle {{ paddle_icon }}

> The PyTorch-TorchScript freeze route does not preserve Hessian output. A
> PyTorch-TorchScript model frozen with `dp --pt freeze` is treated as a
> standard energy model. PyTorch-Exportable cannot construct a Hessian model
> for freezing or inference. The JAX backend can preserve Hessian output in a
> frozen model with `dp --jax freeze --hessian`.

> The PyTorch-TorchScript freeze route cannot output Hessians, and
> PyTorch-Exportable cannot construct a Hessian model for freezing or
> inference. The JAX backend can retain Hessian output in a frozen model with
> `dp --jax freeze --hessian`; see [Freeze a model](../freeze/freeze.md).